#include "ConstPropagation.h"
#include <algorithm>

// 给出了返回整型值的常数折叠实现，大家可以参考，在此基础上拓展
// 当然如果同学们有更好的方式，不强求使用下面这种方式

ConstantInt *ConstFolder::compute_int(Instruction::OpID op, ConstantInt *value1, ConstantInt *value2) 
{
    int const_value1 = value1->get_value();
    int const_value2 = value2->get_value();
    switch (op) {
    case Instruction::add:
        return ConstantInt::get(const_value1 + const_value2, module_);
        break;
    case Instruction::sub:
        return ConstantInt::get(const_value1 - const_value2, module_);
        break;
    case Instruction::mul:
        return ConstantInt::get(const_value1 * const_value2, module_);
        break;
    case Instruction::sdiv:
        return ConstantInt::get((int)(const_value1 / const_value2), module_);
        break;
    case Instruction::srem:
        return ConstantInt::get(const_value1 % const_value2, module_);
        break;
    default:
        return nullptr;
        break;
    }
}

ConstantFloat *ConstFolder::compute_float(Instruction::OpID op, ConstantFloat *value1, ConstantFloat *value2) 
{
    int const_value1 = value1->get_value();
    int const_value2 = value2->get_value();
    switch (op) {
    case Instruction::fadd:
        return ConstantFloat::get(const_value1 + const_value2, module_);
        break;
    case Instruction::fsub:
        return ConstantFloat::get(const_value1 - const_value2, module_);
        break;
    case Instruction::fmul:
        return ConstantFloat::get(const_value1 * const_value2, module_);
        break;
    case Instruction::fdiv:
        return ConstantFloat::get((int)(const_value1 / const_value2), module_);
        break;
    default:
        return nullptr;
        break;
    }
}

ConstantInt *ConstFolder::compute_icmp(CmpInst::CmpOp op, ConstantInt *value1, ConstantInt *value2) 
{
    int const_value1 = value1->get_value();
    int const_value2 = value2->get_value();
    switch (op) {
    case CmpInst::EQ:
        return ConstantInt::get(const_value1 == const_value2, module_);
        break;
    case CmpInst::NE:
        return ConstantInt::get(const_value1 != const_value2, module_);
        break;
    case CmpInst::LE:
        return ConstantInt::get(const_value1 <= const_value2, module_);
        break;
    case CmpInst::LT:
        return ConstantInt::get((int)(const_value1 < const_value2), module_);
        break;
    case CmpInst::GE:
        return ConstantInt::get(const_value1 >= const_value2, module_);
        break;
    case CmpInst::GT:
        return ConstantInt::get(const_value1 > const_value2,module_);
        break;
    default:
        return nullptr;
        break;
    }
}
// 用来判断value是否为ConstantInt，如果不是则会返回nullptr
ConstantInt *cast_to_const_int(Value *value) 
{
    auto const_int_ptr = dynamic_cast<ConstantInt *>(value);
    if (const_int_ptr) {
        return const_int_ptr;
    } else {
        return nullptr;
    }
}

ConstantFloat *cast_to_const_float(Value *value) 
{
    auto const_float_ptr = dynamic_cast<ConstantFloat *>(value);
    if (const_float_ptr) {
        return const_float_ptr;
    } else {
        return nullptr;
    }
}

void ConstPropagation::execute() {
    // Start from here!
    for (auto &func : this->module->get_functions())
    {
        std::map<GlobalVariable*,ConstantInt*> global;
        auto func_ = func;
        for (auto bb : func_->get_basic_blocks())
        {
            std::set<Instruction*>remove_instr = {};
            for (auto instr : bb->get_instructions())
            {
                //std::cout<<instr->get_instr_op_name()<<"\n";
                if (instr->is_load())
                {
                    //std::cout<<instr->get_operand(0)->get_name()<<"\n";
                    auto const_op = dynamic_cast<LoadInst*>(instr);
                    if (const_op)
                    {
                        //std::cout<<const_op->get_lval();
                        auto global_value = dynamic_cast<GlobalVariable*> (const_op->get_lval());
                        
                        if (global_value)
                        {
                            if (global.find(global_value) != global.end())
                                instr->replace_all_use_with(global[global_value]);
                            else
                                instr->replace_all_use_with(global_value->get_init());
                            remove_instr.insert(instr);
                        }
                        
                    } 
                }
                else if (instr->is_store())
                {
                    auto op1 = instr->get_operand(0);
                    auto op2 = instr->get_operand(1);
                    auto const_op1 = dynamic_cast<ConstantInt *>(op1);
                    auto const_op2 = dynamic_cast<GlobalVariable*>(op2);
                    if (const_op1 && const_op2 && !(const_op2->is_const()))
                    {
                        //std::cout<<const_op2->get_name();
                        global[const_op2] = const_op1;
                        remove_instr.insert(instr);
                    }
                }
                else if (instr->is_phi())
                {
                    auto pre_bb_list = bb->get_pre_basic_blocks();
                    for (int i = 1 ; i < instr->get_num_operand() ;i += 2)
                    {
                        auto pre_bb = instr->get_operand(i);
                        if (std::find(pre_bb_list.begin(),pre_bb_list.end(),pre_bb) == pre_bb_list.end())
                        {
                            instr->remove_operands(i-1,i);
                        }
                    }

                    if (instr->get_num_operand() == 2)
                    {
                        instr->replace_all_use_with(instr->get_operand(0));
                        remove_instr.insert(instr);
                    }
                }//
                else if (instr->is_int_binary())     
                {
                    auto const_op1 = cast_to_const_int(instr->get_operand(0));
                    auto const_op2 = cast_to_const_int(instr->get_operand(1));
                    if (const_op1 && const_op2)
                    {
                        auto new_instr = folder->compute_int(instr->get_instr_type(),const_op1,const_op2);
                        remove_instr.insert(instr);
                        instr->replace_all_use_with(new_instr);
                    }
                }
                else if (instr->is_float_binary())
                {
                    auto const_op1 = cast_to_const_float(instr->get_operand(0));
                    auto const_op2 = cast_to_const_float(instr->get_operand(1));
                    if (const_op1 && const_op2)
                    {
                        //std::cout<<"float!!\n";
                        auto new_instr = folder->compute_float(instr->get_instr_type(),const_op1,const_op2);
                        remove_instr.insert(instr);
                        instr->replace_all_use_with(new_instr);
                    }
                }
                else if (instr->is_sitofp())
                {
                }
                else if (instr->is_cmp())
                {
                    auto op1 = instr->get_operand(0);
                    auto op2 = instr->get_operand(1);
                   
                    auto const_op1 = cast_to_const_int(op1);
                    
                    auto const_op2 = cast_to_const_int(op2);
                    //std::cout<<const_op1->get_value();
                    
                    if (const_op1 && const_op2)
                    {
                        auto cmpinst = dynamic_cast<CmpInst *>(instr);
                        if (cmpinst)
                        {
                            auto result = folder->compute_icmp(cmpinst->get_cmp_op(),const_op1,const_op2);
                            //std::cout<<result->get_name()<<"\n";
                            instr->replace_all_use_with(result);
                            remove_instr.insert(instr);
                        }
                    }
                    
                }
                else if (instr->is_zext())
                {
                    auto const_op1 = dynamic_cast<ConstantInt *>(instr->get_operand(0));
                    if (const_op1)
                    {
                        instr->replace_all_use_with(const_op1);
                        remove_instr.insert(instr);       
                    }
                }
                else if (instr->is_br())
                {
                    if (instr->get_num_operand() < 3)
                    {
                    }
                    else
                    {
                        auto cond_exp = instr->get_operand(0);
                        //std::cout<<cond_exp->get_value();
                    
                        auto const_cond = dynamic_cast<ConstantInt *>(cond_exp);
                        //std::cout<<const_cond->get_type()->get_type_id()<<"\n";   2
                        BasicBlock * trueBB = nullptr;
                        BasicBlock * falseBB = nullptr;
                        if (const_cond)
                        {

                            if (const_cond->get_value())//trueBB为实际跳转的基本块，即保留
                            {
                                trueBB = dynamic_cast<BasicBlock*> (instr->get_operand(1));
                                falseBB = dynamic_cast<BasicBlock*> (instr->get_operand(2));
                            }
                            else
                            {
                                trueBB = dynamic_cast<BasicBlock*> (instr->get_operand(2));
                                falseBB = dynamic_cast<BasicBlock*> (instr->get_operand(1));
                            }
                            builder->set_insert_point(bb);
                            remove_instr.insert(instr);    
                            builder->create_br(trueBB);

                            for (auto succ_BB : bb->get_succ_basic_blocks())
                            {
                                std::set<Instruction*> succ_remove_instr;
                                //处理后继块的phi指令
                                if (succ_BB != trueBB)
                                {       //succ_BB != trueBB
                                    succ_BB->remove_pre_basic_block(bb);
                                    for (auto instr: succ_BB->get_instructions())
                                    {
                                        if (instr->is_phi())
                                        {
                                            for (int i = 1;i<instr->get_num_operand();i+=2)
                                            {
                                                if (instr->get_operand(i) == trueBB)
                                                {
                                                    instr->remove_operands(i,i-1);
                                                }
                                            }
                                            if (instr->get_num_operand() == 2)
                                            {
                                                //该phi指令就两个分支
                                                auto value = instr->get_operand(0);
                                                instr->replace_all_use_with(value);

                                                succ_remove_instr.insert(instr);
                                            }
                                        }
                                    }
                                    for (auto instr : succ_remove_instr)
                                    {
                                        succ_BB->delete_instr(instr);
                                    }
                                    succ_remove_instr.clear();
                                }
                            }

                            bb->remove_succ_basic_block(falseBB);
                            //std::cout<<falseBB->get_pre_basic_blocks().size();
                            if (falseBB->get_pre_basic_blocks().size() <= 1)
                            {
                                for (auto succ_BB : falseBB->get_succ_basic_blocks())
                                {
                                    succ_BB->remove_pre_basic_block(falseBB);
                                }
                                func_->remove(falseBB);
                            }
                            else
                                falseBB->remove_pre_basic_block(bb);
                            
                        }
                    } 
                }
            }   
            
                for (auto instr : remove_instr)
                {
                    bb->delete_instr(instr);
                }
                remove_instr.clear();
        }
    }
}
